import asyncio
import openai
import pymysql
import datetime,time, re
from pytz import timezone
from wechaty import (
    Message,
    Wechaty,
    MessageType,
    Friendship,
    UrlLink,
    Contact
)
import os
import sys

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
sys.path.append(os.path.abspath(os.curdir))
# puppet_padlocal_ddebd6ed53a1458b98c9f6dcfeaa2f72  6.6号
# puppet_padlocal_2f41b39641524d9c9beb401a494db04f  10.20号

# 数据库
class DoMysql:
    # 初始化
    def __init__(self):
        # 创建连接
        self.conn = pymysql.Connect(
            host='if.prowealthgroup.cn',
            port=3306,
            user='pwfeishi',
            password='puwei666',
            db='web_bullet',
            charset='utf8',
            cursorclass=pymysql.cursors.DictCursor  # 以字典的形式返回数据
        )
        # 获取游标
        self.cursor = self.conn.cursor()
    # 返回多条数据
    def fetchAll(self, sql):
        self.cursor.execute(sql)
        return self.cursor.fetchall()

    # 返回一条数据
    def fetchOne(self, sql):
        self.cursor.execute(sql)
        return self.cursor.fetchone()

    # 更新数据
    def update(self, sql):
        result = self.cursor.execute(sql)
        return result

    # 关闭连接
    def close(self):
        self.cursor.close()
        self.conn.close()

    # 提交
    def commit(self):
        self.conn.commit()

    # 插入一条数据
    def insert_one(self, sql):
        result = self.cursor.execute(sql)
        return result

    # 插入多条数据
    def insert_many(self, sql, datas):
        result = self.cursor.executemany(sql, datas)
        return result

# 操作
class Act():
    def addslashes(self, s):
        if (isinstance(s, str)):
            d = {"\0": "", "&nbsp;": " ", "\\": "\\\\", "'": "\\'", "\"": "\\\"", "\r": "\\r", "\n": "\\n"}
            for x in d:
                s = s.replace(x, d.get(x))
            return s
            # return "'"+pymysql.escape_string(s)+"'"
        elif s is None:
            return "NULL"
        else:
            return str(s)

class lower_chatGPT:
    # chatgpt预处理
    def __init__(self):
        self.memory_len = 20
        self.memory = {}
        self.apikey = "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"
        self.act = Act()

    # 收到/回复
    async def on_message(self, msg: Message):
        # 如果是自己的信息 或 非文本信息 则跳过
        if msg.is_self() or msg.type() != MessageType.MESSAGE_TYPE_TEXT:
            return
        # ========================== 1. 获取接收信息 ==========================
        # 消息发送者
        talker = msg.talker()
        id = talker.contact_id
        # 发送消息人ID
        send_user_id = talker.get_id()
        # 发送人名称
        send_username = talker.name
        # 接收文本
        text = msg.text()
        # 群聊
        room = msg.room()
        # ========================== 2. 是群聊且没@机器人 跳过  ==========================
        if msg.room() and "@普伟GPT" not in text:
        # if room and "@robot" not in text:
            return
        # ========================== 3. 定义默认数据 ==========================
        text.replace('@普伟GPT', '')
        # text = text.replace('@robot', '')
        mysql = DoMysql()
        # 当前时间
        # now_date = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        tz = timezone('Asia/Shanghai')
        now_date = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        if text == 'test':
            url = 'https://t.1yb.co/m12h'
            url_link = Message.toUrlLink(url)
            # await from_contact.say(url_link)
            return
        # 当前模式（单聊0 群聊1）
        now_mode = 1 if room else 0
        # 用户初始的token
        num = 0
        # 匹配的文章
        articles = []
        # ============= token使用量 =============
        token1 = 0
        token2 = 0
        token3 = 0
        token4 = 0
        token5 = 0
        # ============= 获取提问prompt =============
        # 问题的chatgpt回复
        message1 = ''
        message2 = ''
        # 获取关键词
        prompt2 = '\n请根据句子的字面意思直接提取5个以内的关键词（每个关键词小于4个汉字），当句子字符少于4个字，关键词为原句，每个关键词两边都用英文的双引号包裹，两个词之间以,分割\n输出格式为：\n"xxx","yyy"\n'
        # 获取问题提问语
        prompt3 = ''
        # 获取课程
        prompt4 = ''
        # ============= 获取提问回答 =============
        result_answer = ''
        result_end = ''
        result_lesson = ''
        result_lesson_links = []
        # ========================== 4. 获取发送消息用户的数据 ==========================
        sql1 = "select * from user where wechat_id ='" + send_user_id + "'"
        user_info = mysql.fetchOne(sql1)
        if user_info:
            num = user_info.get('token_num')
            now_mode = 1 if room else int(user_info.get('mode'))
        else:
            await msg.say("用户异常，请联系管理员处理")
            return
        user_id = user_info.get('id')
        # ========================== 5. 判断用户输入特定信息 ==========================
        # ============ 5.1 模式切换 ============
        if not room and text == '原版模式':
            update_mode_sql = "update `user` set `mode`= 0 where id =" + str(user_id)
            result_update_mode = mysql.update(update_mode_sql)
            if result_update_mode:
                mysql.commit()
                mysql.close()
            await msg.say("已切换到原版模式，若您需要普伟客服模式，输入指令“客服模式”，即可完成切换")
            return
        elif not room and text == '客服模式':
            update_mode_sql = "update `user` set `mode`= 1 where id =" + str(user_id)
            result_update_mode = mysql.update(update_mode_sql)
            if result_update_mode:
                mysql.commit()
                mysql.close()
            await msg.say("已切换到客服模式，若您需要chatgpt聊天模式，输入指令“原版模式”即可完成切换")
            return
        try:
        # ============ 5.2 修改用户聊天记录 ============
            check = re.match('^id[::=](\d+)的?聊天记录值改[成为](\d+)$', text)
            if check:
                change_id = check.group(1).strip()
                change_history = check.group(2).strip()
                sql8 = "select * from user where id ='" + change_id + "'"
                sql9 = "update `user` set `history_num`='" + str(change_history) + "' where id =" + str(change_id)
                user_change = mysql.fetchOne(sql8)
                if user_change:
                    change_wechat_id = user_change.get('wechat_id')
                    self.memory[change_wechat_id] = []
                    mysql.update(sql9)
                    mysql.commit()
                    mysql.close()
                    await msg.say("修改成功")
                else:
                    await msg.say("id为" + str(change_id) + "的用户不存在")
                return
        # ========================== 6. 用户模式处理 ==========================
            openai.api_key = self.apikey
            # 是否暂设为普通模式（是否没匹配文章也没匹配课程）
            normal_state = 0
            # ============= 6.1 普伟模式 =============
            if now_mode:
                # ============= 1) 提取问题的关键词 =============
                prompt2 = str(text) + prompt2
                print('------------------')
                print(prompt2)
                response2 = openai.Completion.create(
                    model="text-davinci-003",
                    prompt=prompt2,
                    temperature=0,
                    max_tokens=200,
                    top_p=1,
                    frequency_penalty=0,
                    presence_penalty=0
                )
                result_keys = response2["choices"][0]["text"].strip()
                print(result_keys)
                print('==================')
                token2 = response2['usage']['total_tokens'] * 1
                # ============= 1.1) 判断关键词是否含有"管理" =============
                if "管理" in result_keys:
                    print('------------------')
                    message2 = [{'role': 'system', 'content': '你现在是一名企业咨询师'}, {'role': 'user', 'content': '客户提出问题，请你先用人的思维模式思考并判断客户问题是否想要“企业管理”“企业咨询”“著名企业家/投资人”“自我提升”等咨询相关信息。请回答Yes或No'}]
                    print(message2)
                    response5 = openai.ChatCompletion.create(
                        model="gpt-4",
                        max_tokens=256,
                        top_p=1,
                        temperature=1,
                        frequency_penalty=0,
                        presence_penalty=0,
                        messages=message2
                    )
                    result_manage = response5["choices"][0]["message"]['content'].strip()
                    token5 = response5['usage']['total_tokens'] * 1
                    if "Yes" not in result_manage:
                        result_keys = result_keys.replace('"管理",',"")
                        result_keys = result_keys.replace('"管理"',"")
                    print('==================')
                    print(result_keys)
                # ============= 2) 查找数据库中是否有这关键词 处理成prompt3和prompt4 否则暂用普通模式 =============
                if re.match('^(.+),?$', result_keys):
                    # ============= 2.1) 查找数据库中有匹配文章 =============
                    sql2 = "SELECT b.bullet_id, b.article_url,b.article_short_key,b.article_docs, SUM(k.word_level) AS word_level, b.article_title, b.article_desc FROM `web_project_bullet_key` AS k LEFT JOIN web_project_bullet AS b ON b.bullet_id = k.bullet_id WHERE k.key_word IN (" + str(
                        result_keys.strip()) + ") AND k.is_delete = 0 AND b.article_url IS NOT NULL GROUP BY b.bullet_id ORDER BY word_level DESC"
                    articles = mysql.fetchAll(sql2)
                    if articles and len(articles):
                        # 文章信息（标题，链接，中心思想，文章语录）
                        prompt3 = "客户向你提出问题，请理解问题从问题里提取关键信息并按照示例模板生成回复推文引导语并加上引号内容“我推荐您阅读以下语录，希望能为您提供xxx帮助：”请阅读以下内容\n示例问题：怎么样管理财务团队\n示例回复：如果您想要学会管理财务团队，我推荐您阅读以下语录，希望能为您学会管理财务团队提供帮助：\n问题：" + str(
                            text) + "\n回复："
                    else:
                        normal_state = 1
                    # ============= 2.2) 查找数据库有匹配课程 =============
                    sql3 = "SELECT l.id, SUM(k.word_level) AS word_level, l.lesson_name, l.`desc`, l.lesson_help,l.title,l.short_msg,l.img_url,l.url FROM `lesson_key` AS k LEFT JOIN lesson AS l ON l.id = k.lesson_id WHERE k.key_word IN (" + str(
                        result_keys.strip()) + ") AND k.is_delete = 0 GROUP BY l.id ORDER BY word_level DESC LIMIT 0,2"
                    lessons = mysql.fetchAll(sql3)
                    if lessons:
                        normal_state = 0
                        prompt4 = '客户向你提出问题，你根据问题推荐相关课程并介绍课程和告诉客户学习课程可以带来的成果和帮助，请用企业咨询师的思维模式并用具有亲和力和说服力的语气利用以下的信息整理一段课程推荐引导语，推荐课程名称用《》包裹（每篇课程推荐的间隔带上换行符）\n问题：' + str(
                            text) + "\n"
                        lesson_str = ''
                        for lesson in lessons:
                            lesson_title = lesson['lesson_name']
                            lesson_desc = lesson['desc']
                            lesson_help = lesson['lesson_help']
                            lesson_str = "推荐课程：《" + str(lesson_title) + "》 \n课程介绍：" + str(
                                lesson_desc) + "\n课程帮助：" + str(lesson_help) + "\n"
                            result_lesson_links.append({'title': lesson['title'], 'description': lesson['short_msg'], 'url': lesson['url'], 'thumbnail_url': lesson['img_url']})
                            prompt4 += lesson_str
                        prompt4 += "\n推荐引导语：\n"
                else:
                    return
            # ============= 6.2 普通模式 或 普伟模式（没匹配文章和课程） =============
            if now_mode == 0 or normal_state:
                if len(str(text)) > 900:
                    await msg.say('问题过长')
                    return
                memory_data = {}
                # 判断该用户是否有记录 有的话拼接
                if not room and id in self.memory.keys():
                    temp = self.memory[id].copy()
                    temp.append({"role": "user", "content": str(text)})
                    message1 = temp
                else:
                    message1 = [{'role': 'user', 'content': str(text)}]
        # ========================== 7. 集中处理prompt1 prompt3 prompt4  ==========================
            result_answer = '（省略）'
            if message1:
                print('------------------')
                print(message1)
                response1 = openai.ChatCompletion.create(
                    model="gpt-4",
                    max_tokens=1000,
                    top_p=1,
                    frequency_penalty=0,
                    presence_penalty=0,
                    messages = message1
                )
                result_answer = response1["choices"][0]["message"]['content'].strip()
                print(result_answer)
                token1 = response1['usage']['total_tokens'] * 1
                print(token1)
                print('==================')
            if prompt3 and articles:
                print('------------------')
                print(prompt3)
                response3 = openai.Completion.create(
                    model="text-davinci-003",
                    prompt=prompt3,
                    temperature=0.7,
                    max_tokens=300,
                    top_p=1,
                    frequency_penalty=0,
                    presence_penalty=0
                )
                result_end = response3["choices"][0]["text"].strip()
                i = 0
                for article in articles:
                    if i < 5:
                        result_end += str(i + 1) + ". " + article.get('article_docs') + "\n"
                        i += 1
                    else:
                        break
                print(result_end)
                print('==================')
                token3 = response3['usage']['total_tokens'] * 1
            if prompt4:
                print('------------------')
                print(prompt4)
                response4 = openai.Completion.create(
                    model="text-davinci-003",
                    prompt=prompt4,
                    temperature=0.7,
                    max_tokens=1000,
                    top_p=1,
                    frequency_penalty=0,
                    presence_penalty=0
                )
                result_lesson = response4["choices"][0]["text"].strip()
                print('lesson_ok')
                print('==================')
                token4 = response4['usage']['total_tokens'] * 1
        # ========================== 8. 把问题与回答记录到聊天记录中  ==========================
            if not room:
                # 获取用户默认长度/2
                use_len = int(user_info.get('history_num') / 2) if user_info.get('history_num') != 20 else self.memory_len
                if id not in self.memory.keys():
                    self.memory[id] = []
                if len(self.memory[id]) > use_len:
                    self.memory[id].pop(0)
                    self.memory[id].pop(0)
                    self.memory[id].append({"role": "user", "content": str(text)})
                    self.memory[id].append({"role": "assistant", "content": result_answer})
                else:
                    self.memory[id].append({"role": "user", "content": str(text)})
                    self.memory[id].append({"role": "assistant", "content": result_answer})
        # ========================== 9. 用户token用量计算 ==========================
            # 记录token用量
            insert_user_tokens = []
            # 总支出（RMB）
            total_amount = 0
            # ========== 9.1 普通模式提问（需支付） ==========
            # chatgpt4：$0.03/1000 token => ￥0.21/1000 token
            # chatgpt3.5：$0.002/1000 token => ￥0.014/1000 token
            if token1:
                rmb = round((0.21 * token1 / 1000), 4)
                dollars = round((0.03 * token1 / 1000), 4)
                insert_user_tokens.append([user_id, 2, token1, rmb, dollars, 1, str(now_date)])
                total_amount += rmb
            # ========== 9.2 普伟gpt模式提问（不支付）==========
            # token2,3,4是chatgpt3.5      token5是chatgpt4
            if token2 or token3 or token4 or token5:
                if token5:
                    dollars = round((0.03 * token5 / 1000), 4)
                    rmb = round((0.21 * token5 / 1000), 4)
                    insert_user_tokens.append([user_id, 2, token5, rmb, dollars, 0, str(now_date)])
                dollars = 0
                rmb = 0
                dollars = round((0.002 * (token2 + token3 + token4)) / 1000, 4)
                rmb = round((0.014 * (token2 + token3 + token4)) / 1000, 4)
                insert_user_tokens.append([user_id, 1, (token2 + token3 + token4), rmb, dollars, 0, str(now_date)])
        # ========================== 10. 保存到数据库 ==========================
            # 用户使用token情况
            sql7 = "insert into `user_token`(`user_id`,`token_type`,`token_num`,`token_amount`,`usa_amount`, `is_pay`, `create_time`) values(%s,%s,%s,%s,%s,%s,%s)"
            mysql.insert_many(sql7, insert_user_tokens)
            # 用户余额，名称修改
            # 收费按成本的一半 total_amount * 1.5
            remaining_sum = round(float(user_info.get('remaining_sum')) - float(total_amount) * 1.5, 4)
            if remaining_sum  < 0.01:
                await msg.say('您的余额不足，请充值！')
                return
            sql10 = "update `user` set `wechat_name`='" + str(send_username) +"', `remaining_sum`= '" + str(remaining_sum) + "', `last_update_time`='" + now_date +"' where id =" + str(user_id)
            mysql.update(sql10)
            mysql.commit()
            mysql.close()
            if result_answer != '（省略）':
                await msg.say(result_answer)
            if result_end:
                await msg.say(result_end[:-1])
            if result_lesson:
                await msg.say(result_lesson)
                for link in result_lesson_links:
                    if link.get('url'):
                        urlLink = UrlLink.create(
                            description = link.get('description') if link.get('description') else 'error',
                            thumbnail_url = link.get('thumbnail_url') if link.get('thumbnail_url') else 'http://www.test.com',
                            title = link.get('title') if link.get('title') else 'error',
                            url = link.get('url') if link.get('url') else 'http://www.test.com',
                        )
                        print(urlLink)
                        await msg.say(urlLink)
        except Exception as e:
            self.memory[id] = []
            await msg.say("ERROR! ")
        time.sleep(0.1)

    # 好友添加
    async def on_friendship(self, friendship: Friendship):
        # 当前时间
        tz = timezone('Asia/Shanghai')
        now_date = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        wechatId = friendship.contact().get_id()
        mysql = DoMysql()
        # 获取接收者是否用户
        sql5 = "select * from user where wechat_id ='" + wechatId + "'"
        user_info = mysql.fetchOne(sql5)
        if user_info:
            pass
        else:
            try:
                sql6 = "INSERT INTO `user` (`wechat_id`, `wechat_name`, `mode`, `status`, `insert_time`, `last_update_time`) VALUES ('" + wechatId + "', '', 0, '1', '"+ now_date + "', '" + now_date + "');"
                mysql.insert_one(sql6)
                mysql.commit()
                mysql.close()
            except:
                print('error')

async def main():
    # os.environ['WECHATY_PUPPET_SERVICE_TOKEN'] = 'insecure_f47ea95c-01a0-4043-ac24-64c65c551207'
    os.environ['WECHATY_PUPPET_SERVICE_TOKEN'] = 'insecure_f47ea95c-01a0-4043-ac24-64c65c558c9f'
    os.environ['WECHATY_PUPPET'] = 'wechaty-puppet-padlocal'
    os.environ['WECHATY_PUPPET_SERVICE_ENDPOINT'] = '127.0.0.1:8089'
    wechat = Wechaty()
    Bot = lower_chatGPT()
    wechat.on('message', Bot.on_message)
    wechat.on('friendship', Bot.on_friendship)

    await wechat.start()

if __name__ == "__main__":
    asyncio.get_event_loop().run_until_complete(main())